#!/usr/bin/env python

from struct import unpack, unpack_from, pack
from functools import partial

def loads(s, strict=True, enable_blobs=False):
    """
    Decode BJSOS stream

    If `enable_blobs` is True, support decoding extended BJSOS w/
    BLOBs. BLOBs are mapped to bytearrays.
    """

    val, rest = loads2(s, strict=strict, enable_blobs=enable_blobs)
    assert rest == "", "trailing garbage"
    return val

def loads2(s, strict=True, enable_blobs=False):
    """
    Variant of loads() which parses the first BJSOS value and returns
    the remaining unparsed string
    """

    assert isinstance(s, bytes)

    val, rest = _decode(s, enable_blobs=enable_blobs)
    if strict:
        assert isinstance(val, (list,dict)), "top-level structure must be object or array in strict mode"

    return val, rest

def dumps(v, enable_blobs=False):
    """
    Serialize value to BJSOS

    If `enable_blobs` is True, encode bytearrays as BJSOS BLOBs
    (unless used as key in a dict); otherwise, bytearrays are
    interpreted as 7-bit ASCII strings and encoded as such.
    """

    dumps_ = partial(dumps, enable_blobs=enable_blobs)

    if v is None:    return b'\x40'
    elif v is False: return b'\x42'
    elif v is True:  return b'\x43'

    elif isinstance(v, (list,tuple)):
        cnts, cntsz = encodeWord(len(v))
        data = cnts + "".join(map(dumps_, v))
        lens, lensz = encodeWord(len(data))
        return pack("!B", 0x70 | cntsz << 2 | lensz) + lens + data

    elif isinstance(v, dict):
        cnts, cntsz = encodeWord(len(v))
        # N.B.: keys are never encoded as BLOBs
        data = cnts + "".join(dumps(k1) + dumps_(v1) for k1,v1 in v.items())
        lens, lensz = encodeWord(len(data))
        return pack("!B", 0x50 | cntsz << 2 | lensz) + lens + data

    elif isinstance(v, unicode):
        bs = v.encode('utf8')
        bsl = len(bs)
        if bsl < 0x40: # small-string
            return pack("!B", 0x80 | bsl) + bs
        else:
            lens, lensz = encodeWord(bsl)
            return pack("!B", 0x4c | lensz) + lens + bs

    elif isinstance(v, bytes):
        if enable_blobs:
            lens, lensz = encodeWord(len(v))
            return pack("!B", 0x48 | lensz) + lens + v
        else:
            return dumps(v.decode('ascii'))

    elif isinstance(v, (int,long)):
        if -0x40 <= v < 0x40:
            return pack("!b", v)
        elif -0x80 <= v < 0x80:
            return pack("!bb", 0x44, v)
        elif -0x8000 <= v < 0x8000:
            return pack("!bh", 0x45, v)
        elif -0x80000000 <= v < 0x80000000:
            return pack("!bi", 0x46, v)
        elif -0x8000000000000000 <= v < 0x8000000000000000:
            return pack("!bq", 0x47, v)
        else:
            raise RuntimeError("cannot serialize large {}".format(type(v)))

    else: raise RuntimeError("cannot serialize {}".format(type(v)))

def encodeWord(v):
    assert v >= 0
    if v <= 0xff:
        return pack("!B", v), 0
    elif v <= 0xffff:
        return pack("!H", v), 1
    elif v <= 0xffffffff:
        return pack("!I", v), 2
    elif v <= 0xffffffffffffffff:
        return pack("!Q", v), 3
    else:
        raise RuntimeError("encodeWord")

def decode_checklen(s, enable_blobs=False):
    """
    Validate completeness of top-level structure in bytestring 's'

    This function is useful for deciding if a BJSOS value in a receive
    buffer is complete, or if additional data needs to be read from
    the socket/file.

    Returns 0 if bytestring exactly as long as the data structure.
    Returns +n if at least n bytes are missing.
    Returns -n if buffer contains n bytes of extraneous data.

    Invariant(s):
     - forall v . decode_checklen(dumps(v)) == 0
    """

    if s == "":
        return 1 # we need at least a tag-byte

    t = ord(s[0])
    l1 = len(s) - 1

    if t <= 0x3f or 0xc0 <= t: # small-int
        return 0 - l1

    elif t >= 0x80: # short-string
        l = t & 0x3f
        return l - l1

    elif t == 0x40: return 0 - l1 # null
    elif t == 0x41: return 8 - l1 # double
    elif t == 0x42: return 0 - l1 # false
    elif t == 0x43: return 0 - l1 # true
    elif t <= 0x47: # int{8,16,32,64}
        szpfx = 1 << (t & 0x3)
        return szpfx - l1
    elif 0x60 <= t and t <= 0x6f: pass # invalid tags
    elif t <= 0x7f: # BLOB/string/object/array
        if t <= 0x4b and not enable_blobs:
            raise RuntimeError("BLOB type encountered but enable_blobs=False")
        szpfx = 1 << (t & 0x3)
        if szpfx > l1:
            return szpfx - l1
        else:
            l, s_ = _decodeWord(t, s[1:])
            return l - len(s_)

    raise RuntimeError("invalid BJSOS tag (0x{0:x})".format(t))


# recursive parsers
# warning: extremely inefficient
# invalid BJSOS serialization may trigger index-out-of-range exceptions
def _decode(s, enable_blobs=False):
    t = ord(s[0])

    if t <= 0x3f or 0xc0 <= t: # small-int
        return (unpack_from("!b",s)[0], s[1:])

    elif t >= 0x80: # short-string
        l = t & 0x3f
        assert len(s) >= 1+l
        return (s[1:1+l].decode('utf8'), s[1+l:])

    elif t == 0x40: return None, s[1:] # null
    elif t == 0x41: return unpack("!d",s[1:9])[0],s[9:]
    elif t == 0x42: return False, s[1:]
    elif t == 0x43: return True, s[1:]
    elif t <= 0x47: return _decodeInt(t, s[1:])
    elif t <= 0x4b: # BLOBs
        if not enable_blobs:
            raise RuntimeError("BLOB type encountered but enable_blobs=False")

        l, s_ = _decodeWord(t, s[1:])
        assert len(s_) >= l
        return s_[:l], s_[l:]

    elif t <= 0x4f: # string
        l, s_ = _decodeWord(t, s[1:])
        assert len(s_) >= l
        return s_[:l].decode('utf8'), s_[l:]

    elif t <= 0x5f: # object
        # common code for object/array
        l, s_ = _decodeWord(t, s[1:])
        assert len(s_) >= l
        s_, rest = s_[:l], s_[l:]
        cnt, s_ = _decodeWord(t>>2, s_)

        # object-specific code
        obj = []
        for _ in xrange(cnt):
            k, s_ = _decode(s_) # BLOBs not allowed in dict keys
            v, s_ = _decode(s_, enable_blobs)
            obj.append((k,v))
        assert s_ == ""
        return dict(obj), rest

    elif t <= 0x6f: pass

    elif t <= 0x7f: # array
        # common code for object/array
        l, s_ = _decodeWord(t, s[1:])
        assert len(s_) >= l
        s_, rest = s_[:l], s_[l:]
        cnt, s_ = _decodeWord(t>>2, s_)

        # array-specific code
        arr = []
        for _ in xrange(cnt):
            v, s_ = _decode(s_, enable_blobs)
            arr.append(v)
        assert s_ == ""
        return arr, rest

    raise RuntimeError("invalid BJSOS tag (0x{0:x})".format(t))

def _decodeInt(t, s):
    t &= 0x3
    if   t == 0: v = unpack("!b", s[:1]), s[1:]
    elif t == 1: v = unpack("!h", s[:2]), s[2:]
    elif t == 2: v = unpack("!i", s[:4]), s[4:]
    elif t == 3: v = unpack("!q", s[:8]), s[8:]
    else: raise RuntimeError("internal error")
    (v,), r = v
    return v,r

def _decodeWord(t, s):
    t &= 0x3
    if   t == 0: v = unpack("!B", s[:1]), s[1:]
    elif t == 1: v = unpack("!H", s[:2]), s[2:]
    elif t == 2: v = unpack("!I", s[:4]), s[4:]
    elif t == 3: v = unpack("!Q", s[:8]), s[8:]
    else: raise RuntimeError("internal error")
    (v,), r = v
    return v,r

if __name__ == "__main__":
    import sys
    import json

    mode = 'bjsos2json'
    if sys.argv[1:] == []:
        mode = 'bjsos2json'
    elif sys.argv[1:] == ["bjsos2json"]:
        mode = 'bjsos2json'
    elif sys.argv[1:] == ["json2bjsos"]:
        mode = 'json2bjsos'
    else:
        print "usage: {} [bjsos2json|json2bjsos]".format(sys.argv[0])
        sys.exit(1)

    if mode == 'bjsos2json':
        sys.stdout.write(json.dumps(loads(sys.stdin.read()))+'\n')
    elif mode == 'json2bjsos':
        sys.stdout.write(dumps(json.loads(sys.stdin.read())))
